import akshare as ak
import os
import pickle
from tqdm import tqdm
import json 
from datetime import datetime
import argparse

def get_code_market(code:str):
    if code.startswith("6"):
        return "sh"
    else:
        return "sz"

def get_code_fund_flow(code):
    market = get_code_market(code)
    return ak.stock_individual_fund_flow(code, market)

def get_data(data_func, cacheFile='stock_fund_flow.data', ):
    if not os.path.exists(cacheFile):
        data = data_func()
        with open(cacheFile, 'wb') as f:
            pickle.dump(data, f)
    else:
        with open(cacheFile, 'rb') as f:
            data = pickle.load(f)
    return data

def get_fund_flows():
    # 获取 A股所有股票代码
    codes = []
    sh_codes = ak.stock_info_sh_name_code()
    for _, code in sh_codes.iterrows():
        codes.append((code["证券代码"], code["证券简称"]))

    sz_codes = ak.stock_info_sz_name_code()
    for _, code in sz_codes.iterrows():
        codes.append((code["A股代码"], code["A股简称"]))
    # cy_codes = ak.stock_cy_a_spot_em()
    # for _, code in cy_codes.iterrows():
    #     codes.append((code["代码"], code["名称"]))

    fundFlows = []
    for (code, name) in tqdm(codes, desc="process"):
        try:
            codeFundFlow = get_code_fund_flow(code)
            fundFlows.append({"code":code, "name":name, "flows":codeFundFlow })
        except Exception as e:
            print(f"{code} {name} err {e}")
            continue
    return fundFlows


def get_day_flow_in_stocks(days):
    # 定义选股条件：连续三日主力流入且小单流出
    # 选出符合条件的股票
    selected_stocks = []
    codesFundFlow = get_data(lambda: get_fund_flows())
    for fund in codesFundFlow:
        try:
            code=fund['code']
            name=fund['name']
            flows=fund['flows']
            selected = True
            zf = []
            if len(flows) < days:
                # 新股天数不足
                continue
            for day in range(-days, 0):
                flow = flows.iloc[day]
                zf.append(format_flow(flow))
                if not (flow['主力净流入-净额'] > 0 and (flow['中单净流入-净额'] + flow['小单净流入-净额']) <= 0 and flow['小单净流入-净额'] < 0):
                    selected = False
                    break
            if selected:
                selected_stocks.append((code, name, zf))
        except Exception as e:
            print(f"{code} err {e}")
            continue
    return selected_stocks

def combo_flow_in(params):
    parser = argparse.ArgumentParser() 
    parser.add_argument("--days", dest="days", help="连续流入天数", type=int)
    args = parser.parse_args(params)
    days = args.days
    print(f'select stock with flow in {days} days')
    selected_stocks = get_day_flow_in_stocks(days)
    # 打印选出的股票代码和涨幅
    filename = f"flow_in_{days}_stock_{datetime.now().strftime('%Y%m%d')}.txt"
    save_result(filename, selected_stocks)

def stock_formater(stocks):
    sep = '\n\t'
    return [f"{code} {name}{sep}{sep.join([json.dumps(d) for d in data])}" for (code, name, data) in stocks]

def raw_formater(datas):
    return [json.dumps(d, ensure_ascii=False) for d in datas]

def save_result(filename, datas, header = "", is_print=False, formater=raw_formater):
    if is_print:
        print(datas)
    with open(filename, "w") as f:
        if header != "":
            f.write(header + '\n')
        f.writelines('\n'.join(formater(datas))) 

def get_day_flow_out_and_rise_stocks(days):
    selected_stocks = []
    codesFundFlow = get_data(lambda: get_fund_flows())
    for fund in tqdm(codesFundFlow, desc="process"):
        try:
            code=fund['code']
            name=fund['name']
            flows=fund['flows']
            selected = True
            zf = []
            if len(flows) < days:
                # 新股天数不足
                continue
            for day in range(-days, 0):
                flow = flows.iloc[day]
                zf.append(format_flow(flow))
                if not (flow['主力净流入-净额'] < -500*10000 and flow['涨跌幅'] > 0):
                    selected = False
                    break
            if selected:
                selected_stocks.append((code, name, zf))
        except Exception as e:
            print(f"{code} err {e} ")
            break
    return selected_stocks

def flow_out_rise(params):
    parser = argparse.ArgumentParser() 
    parser.add_argument("--days", dest="days", help="连续主力流出天数且上涨", type=int)
    args = parser.parse_args(params)
    print(f'select stock with flow out and rise ')
    days = args.days
    selected_stocks = get_day_flow_out_and_rise_stocks(days)
    print(selected_stocks)
    filename = f"flow_out_and_rise_{days}_stock_{datetime.now().strftime('%Y%m%d')}.txt"
    save_result(filename, selected_stocks)


def get_day_flow_out_and_rise_rate(days, next_days):
    selected_stocks = []
    codesFundFlow = get_data(lambda: get_fund_flows())
    for fund in tqdm(codesFundFlow, desc="progress"):
        try:
            stock_hist = []
            code=fund['code']
            name=fund['name']
            flows=fund['flows']
            if len(flows) < days:
                # 新股天数不足
                continue
            for day_index in range(days, len(flows) - next_days + 1):
                zf = []
                out_and_rise = True
                for day in range(-days + day_index, day_index):
                    flow = flows.iloc[day]
                    zf.append(format_flow(flow))
                    if not (flow['主力净流入-净额'] < -1000*10000 and flow['涨跌幅'] > 0):
                        out_and_rise = False
                        break
                if out_and_rise:
                    stock_hist.append((zf, [format_flow(flows.iloc[day_index + i]) for i in range(next_days)]))
            if stock_hist != []:
                selected_stocks.append((code, name, stock_hist))
        except Exception as e:
            print(f"{code} err {e} ")
            break
    return selected_stocks

def format_flow(flow):
    return {"date":flow['日期'], "rise":flow['涨跌幅'], "flow":flow['主力净流入-净额'], "price":flow['收盘价']}

def flow_out_rise_rate(params):
    parser = argparse.ArgumentParser() 
    parser.add_argument("--days", dest="days", help="连续主力流出天数且上涨", type=int)
    parser.add_argument("--next", dest="next", help="之后几天的涨幅", type=int)
    args = parser.parse_args(params)
    days = args.days
    next = args.next
    print(f'select stock with flow out and rise rate {days} next {next}')
    selected_stocks = get_day_flow_out_and_rise_rate(days, next)
    win_result = []
    for (code, name, stock_hist) in selected_stocks:
        code_total = 0
        code_win = 0
        for (zf, next) in stock_hist:
            next_rise = sum(flow['rise'] for flow in next)
            total_rise = sum(flow['rise'] for flow in zf) + next_rise
            code_total += 1
            if next_rise > 0:
                code_win += 1
        win_result.append((code, name, {'win':code_win, 'total':code_total, 'win_rate':code_win/code_total, 'rise':next_rise}))
    win_result.sort(key=lambda x:x[2]['win_rate']*10000000 + x[2]['total'], reverse=True)
    
    filename = f"flow_out_and_rise_rate_{days}_stock_{datetime.now().strftime('%Y%m%d')}.txt"
    save_result(filename, selected_stocks, is_print=False, formater=stock_formater)
    win_filename = f"win_result_{days}_stock_{datetime.now().strftime('%Y%m%d')}.txt"
    win = sum([x[2]['win'] for x in win_result])
    total = sum([x[2]['total'] for x in win_result])
    save_result(win_filename, win_result, is_print=False, header=f"win = {win}, total={total} rate {win/total}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser() 
    parser.add_argument(dest="type", help="策略类型", choices=("flow-in", "flow-out", "flow-out-rate"), type=str)
    (args, handler_args) = parser.parse_known_args()
    handlers = {
        "flow-in": combo_flow_in,
        "flow-out": flow_out_rise,
        "flow-out-rate": flow_out_rise_rate,
    }
    handlers[args.type](handler_args)
    pass
